package com.ly.demo.service.impl;

import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.time.Duration;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

@Service
public class SseService {

    private final Map<String, SseEmitter> emitters = new ConcurrentHashMap<>();
    private final ExecutorService executorService = Executors.newCachedThreadPool();

    private final static Long CHAT_TIMEOUT = Duration.ofMinutes(1).toMillis();

    /**
     * 创建一个新的 SseEmitter 并存储在 emitters Map 中
     *
     * @param unionKey
     * @return
     */
    public SseEmitter createEmitter(String unionKey) {
        SseEmitter emitter = new SseEmitter(CHAT_TIMEOUT);
        emitter.onCompletion(() -> emitters.remove(unionKey));
        emitter.onTimeout(() -> emitters.remove(unionKey));
        emitter.onError((ex) -> {
            emitters.remove(unionKey);
            ex.printStackTrace();
        });
        try {
            emitter.send(SseEmitter.event().name("connect").data("hello world,该链接在未发送消息"+CHAT_TIMEOUT/1000+"秒后断开"));
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        emitters.put(unionKey, emitter);
        return emitter;
    }

    /**
     * 发送事件给特定的客户端
     *
     * @param unionKey
     * @param event
     * @param data
     */
    public void sendEvent(String unionKey, String event, Object data) {
        SseEmitter emitter = emitters.get(unionKey);
        if (emitter != null) {
            executorService.execute(() -> {
                try {
                    emitter.send(SseEmitter.event().name(event).data(data));
                } catch (IOException e) {
                    emitters.remove(unionKey);
                    emitter.completeWithError(e);
                }
            });
        }
    }

    /**
     * 关闭特定客户端的连接
     *
     * @param unionKey
     */
    public void complete(String unionKey) {
        SseEmitter emitter = emitters.get(unionKey);
        if (emitter != null) {
            emitter.complete();
            emitters.remove(unionKey);
        }
    }

    /**
     * 关闭所有客户端的连接
     */
    public void completeAll() {
        emitters.forEach((unionKey, emitter) -> emitter.complete());
        emitters.clear();
    }
}
